#include <stdio.h>
#include <stdlib.h>
#include "cutil.h"

// Dimensions of the image
#define IMG_WIDTH 512
#define IMG_HEIGHT 512

// Samples per pixel (super sampling)
#define SAMPLESZ   3

#define WIDTH  (IMG_WIDTH * SAMPLESZ)
#define HEIGHT (IMG_HEIGHT * SAMPLESZ)
#define NDATA  (5) // R,G,B,X,Y

// GRIDSZ * BLOCKSZ <= 64k because Meresenne Twister can only do 64k threads.
#define BLOCKSZ 512 // 2^9
#define GRIDSZ  128 // 2^7
#define GRIDSZ_IFS (NUM_COORDS / BLOCKSZ)
#define AUTO_ITERATIONS 128 // 2^7 ~= 100
#define NUM_COORDS 512 * 1024 // 2^19 

#define MAX_FUNCTIONS 10

const unsigned long long MT_RNG_COUNT = (BLOCKSZ * GRIDSZ);
const unsigned long long RAND_N = NUM_COORDS * AUTO_ITERATIONS; // See README
const unsigned long long N_PER_RNG = RAND_N / MT_RNG_COUNT;
typedef enum
{
    RED = 0,
    GREEN,
    BLUE,
    NUM_COLORS
} colors;
int numFunctions;
float coeffMatrix[MAX_FUNCTIONS][6];
int functionType[MAX_FUNCTIONS];
float variationWeights[MAX_FUNCTIONS] =
{
    2.0  /* Linear */,
    1.0, /* Sinusodial */
    1.0, /* Spherical */
    0.8, /* Swirl */
    0.7, /* Heart */
    0.7, /* Handkerchief */
    0.9, /* Fisheye */
    0.2, /* NoiseSquare */
    0.3  /* NoiseCircle */
};
float variationColors[MAX_FUNCTIONS][NUM_COLORS];
// Why does there need to be an additional probability when a) the functions
// are already chosen according to the weights given in the settings, and b)
// the function indices are randomly chosen?
float probability[MAX_FUNCTIONS];

void autoFocus(float *xmax, float *xmin, float *ymax, float *ymin);
float randFloat(float min, float max);
int randInt(int min, int max);
void generateRandoms(float xmax, float xmin, float ymax, float ymin);
void initIFS(void);
void generateRawImageData(float xmin, float xmax, float ymin, float ymax,
                          int (*img)[HEIGHT][NDATA], int imageSize);

float *d_Rand;
float *d_Probability;
float (*d_Coefficients)[6];
int (*d_Img)[HEIGHT][NDATA];

__host__ __device__ int weightedRandom(float *weights, int count, float randomNum);

#include "cuRandom.cu"

// How To Parallelize An IFS
//
// If an IFS is "well-formed", applying the IFS to most starting points
// whose coordinates lie inside the "viewable region" will yield another
// point inside the viewable region. So no matter what point you start
// with, the next point is probably going to be inside this region.
// 
// As the algorithm continues to operate (in a single thread, let's say),
// the points generated in each pass tend to be constrained to certain
// regions of the plane by the algorithm. In theory, what we want to do
// is take each point in the viewable region and run it through the IFS
// an infinite number of times (discarding the results) so that we're
// sure it lies on the plane of the IFS. Now that we have an infinite
// number of points which have a 2D distribution determined by the IFS,
// we run each point through the IFS once and record the result.
// Ultimately, what is getting recorded is the result of running all the
// points that are *already* constrained to the range of the IFS back
// through the IFS.
// 
// In practice, what we do is to essentially take a random sample of that
// last step. Given a single point that is constrained to the range of
// the IFS, run it through the IFS and record the result. Repeat a bunch
// of times.
// 
// So how do we get points that are on the IFS to begin with? As it turns
// out, all we need to do is take a random point in the viewable region
// and run it through the IFS about a dozen times to achieve a point
// that's sufficiently close to being constrained to the range of the
// infinitely-iterated IFS to be visually indistinguishable.
// Floating-point errors prevent it from getting any more precise with
// continued iteration.
// 
// Both CPU and GPU threads follow that pattern: pick a different random
// point in each thread, and iterate it a dozen or so times before
// starting to accumulate information into the buffer. From a statistical
// standpoint, there's no real difference between a point that's been
// through the IFS a dozen or a million times, so it doesn't matter if a
// thread only produces a thousand points as long as it warms up its
// starting point first.

int main()
{
    // These are the dimensions of the space from
    // which random points are chosen.
    float xmin = -2;
    float xmax = 2;
    float ymin = -2;
    float ymax = 2;

    int imageSize = sizeof(int) * WIDTH * HEIGHT * NDATA;

    // Seed the CPU random number generator
    srand(42);

    // Allocate a buffer to hold image data
    int (*rawData)[HEIGHT][NDATA] = 
        (int (*)[HEIGHT][NDATA])malloc(imageSize);

    memset(rawData, 0, imageSize);

    initIFS();

    autoFocus(&xmax, &xmin, &ymax, &ymin);

    generateRawImageData(xmin, xmax, ymin, ymax, rawData, imageSize);

    free(rawData);
    CUDA_SAFE_CALL(cudaFree(d_Rand));
    CUDA_SAFE_CALL(cudaFree(d_Img));

    return 0;
}

void
generateRawImageData
(
    float xmin,
    float xmax,
    float ymin,
    float ymax,
    int (*img)[HEIGHT][NDATA],
    int imageSize
)
{
#if 0
    dim3 gridDim;
    gridDim.x = WIDTH / BLOCKSZ;
    gridDim.y = HEIGHT;

    CUDA_SAFE_CALL(cudaMalloc((void **)&d_Img, imageSize));

    generateImageData<<<gridDim, BLOCKSZ>>>(d_Rand,
                                           d_Img,
                                           WIDTH,
                                           HEIGHT,
                                           NUM_COORDS,
                                           xmax - xmin,
                                           xmin,
                                           ymax - ymin,
                                           ymin);

    CUDA_SAFE_CALL(cudaThreadSynchronize());
    CUDA_SAFE_CALL(cudaMemcpy(img, d_Img, imageSize, cudaMemcpyDeviceToHost));
#endif
#if 1
    CUDA_SAFE_CALL(cudaMalloc((void **)&d_Img, imageSize));

    generateImageDataAtomic<<<GRIDSZ_IFS, BLOCKSZ>>>(d_Rand,
                                           d_Img,
                                           WIDTH,
                                           HEIGHT,
                                           NUM_COORDS,
                                           xmax - xmin,
                                           xmin,
                                           ymax - ymin,
                                           ymin);

    CUDA_SAFE_CALL(cudaThreadSynchronize());
    CUDA_SAFE_CALL(cudaMemcpy(img, d_Img, imageSize, cudaMemcpyDeviceToHost));
#endif
}


// Autofocus: we can further shorten the view region to only include the
// most populated region by calculating the standard deviation. A large
// standard deviation indicates that the data points are far from the mean
// and a small standard deviation indicates that they are clustered closely
// around the mean. By choosing the region to be around mean +/- 2stdd, we
// narrow down the region.
void autoFocus
(
    float *xmax,
    float *xmin,
    float *ymax,
    float *ymin
)
{
#if 1
    float stddX, stddY;
    float meanX, meanY;
    float *d_stddX, *d_stddY;
    float *d_meanX, *d_meanY;

    // Generate a set of RAND_N values, which will be later used to form
    // NUM_COORDS co-ordinates. These RAND_N values will be between (0, 1]
    // These values are stored in the d_Rand array.
    printf("Generating random numbers on GPU...\n\n");
    generateRandoms(*xmax, *xmin, *ymax, *ymin);

    // Next, form (x,y) points with these randoms. Use coalesced reads to read
    // from the d_Rand array.
 
    // Run them through the IFS multiple times and store the resulting
    // coordinates.
    printf("Running random numbers through the IFS...\n\n");
    iterateIFS<<<GRIDSZ_IFS, BLOCKSZ>>>(d_Rand,
                                        d_Probability,
                                        d_Coefficients,
                                        numFunctions,
                                        *xmax - *xmin,
                                        *xmin,
                                        *ymax - *ymin,
                                        *ymin);

    CUDA_SAFE_CALL(cudaThreadSynchronize());

    // Using reduction calculate the mean of the values inside d_Rand along
    // the X axis and the Y axis. d_Rand[0..(RAND_N/2 - 1)] will be used
    // the X-mean and d_Rand[RAND_N/2 - (RAND_N - 1)] will be used for the
    // Y-mean.
    CUDA_SAFE_CALL(cudaMalloc((void **)&d_meanX, sizeof(float)));
    CUDA_SAFE_CALL(cudaMalloc((void **)&d_meanY, sizeof(float)));

    printf("Calculating mean...\n\n");
    reduceArrayToMean<<<((2 * NUM_COORDS) / BLOCKSZ), BLOCKSZ>>>(d_Rand,
                                                                 NUM_COORDS,
                                                                 d_meanX,
                                                                 d_meanY);
    CUDA_SAFE_CALL(cudaThreadSynchronize());
    CUDA_SAFE_CALL(cudaMemcpy(&meanX, d_meanX, sizeof(float), cudaMemcpyDeviceToHost));
    CUDA_SAFE_CALL(cudaMemcpy(&meanY, d_meanY, sizeof(float), cudaMemcpyDeviceToHost));
    
    CUDA_SAFE_CALL(cudaMalloc((void **)&d_stddX, sizeof(float)));
    CUDA_SAFE_CALL(cudaMalloc((void **)&d_stddY, sizeof(float)));

    printf("Calculating standard deviation...\n\n");
    calculateStdd<<<((2 * NUM_COORDS) / BLOCKSZ), BLOCKSZ>>>(d_Rand,
                                                            NUM_COORDS,
                                                            d_stddX,
                                                            d_stddY,
                                                            d_meanX,
                                                            d_meanY);
    CUDA_SAFE_CALL(cudaThreadSynchronize());

    CUDA_SAFE_CALL(cudaMemcpy(&stddX, d_stddX, sizeof(float), cudaMemcpyDeviceToHost));
    CUDA_SAFE_CALL(cudaMemcpy(&stddY, d_stddY, sizeof(float), cudaMemcpyDeviceToHost));

    printf("Standard deviation along x: %f, y: %f\n", stddX, stddY);
    printf("Mean along x: %f, y: %f\n", meanX, meanY);

    *xmin = meanX - 2 * stddX; 
    *xmax = meanX + 2 * stddX; 
    *ymin = meanY - 2 * stddY; 
    *ymax = meanY + 2 * stddY; 

    printf("New boundaries: xmin = %f, xmax = %f, ymin = %f, ymax = %f\n",
            *xmin, *xmax, *ymin, *ymax);

#endif
#if 0
    float testArray[] = {1, 0, 1, 0, 1, 0, 1, 0, /*X<-->Y*/ 3, 3, 0, 0, 3, 3, 0, 0};
    float *d_test;
    float stddX, stddY;
    float *d_stddX, *d_stddY;

    printf("Before: ");
    for (int i = 0; i < 16; i++)
    {
        printf("%f ", testArray[i]);
    }
    printf("\n");

    CUDA_SAFE_CALL(cudaMalloc((void **)&d_stddX, sizeof(float)));
    CUDA_SAFE_CALL(cudaMalloc((void **)&d_stddY, sizeof(float)));

    CUDA_SAFE_CALL(cudaMalloc((void **)&d_test, sizeof(testArray)));
    CUDA_SAFE_CALL(cudaMemcpy(d_test, testArray, sizeof(testArray), cudaMemcpyHostToDevice));

    reduceArrayToMean<<<1, BLOCKSZ>>>(d_test, 8, d_stddX, d_stddY);
    CUDA_SAFE_CALL(cudaThreadSynchronize());
    
    // Test
    CUDA_SAFE_CALL(cudaMemcpy(&stddX, d_stddX, sizeof(float), cudaMemcpyDeviceToHost));
    CUDA_SAFE_CALL(cudaMemcpy(&stddY, d_stddY, sizeof(float), cudaMemcpyDeviceToHost));

    printf("meanX: %f, meanY: %f\n", stddX, stddY);
    calculatStdd<<<1, BLOCKSZ>>>(d_test, 8, d_stddX, d_stddY);
    CUDA_SAFE_CALL(cudaThreadSynchronize());

    CUDA_SAFE_CALL(cudaMemcpy(testArray, d_test, sizeof(testArray), cudaMemcpyDeviceToHost));
    printf("After: ");
    for (int i = 0; i < 16; i++)
    {
        printf("%f ", testArray[i]);
    }
    printf("\n");

    CUDA_SAFE_CALL(cudaMemcpy(&stddX, d_stddX, sizeof(float), cudaMemcpyDeviceToHost));
    CUDA_SAFE_CALL(cudaMemcpy(&stddY, d_stddY, sizeof(float), cudaMemcpyDeviceToHost));

    printf("stddX: %f, stddY: %f\n", stddX, stddY);
#endif

    // Once the mean is calculated, use reduction again to calculate the
    // variance along the X-axis and Y-axis. The calculate the standard
    // deviation. Once that is done, the new xmax and ymax will be
    // (mean +/ stdd)
}

void generateRandoms 
(
    float xmax,
    float xmin,
    float ymax,
    float ymin
)
{
    // Storage for RAND_N randoms on the GPU
    CUDA_SAFE_CALL(cudaMalloc((void **)&d_Rand, RAND_N * sizeof(float)));

    // Initialize and seed Mersenne Twister RNG
    loadMTGPU("./MersenneTwister.dat");
    seedMTGPU(42);

    int numIterations = 100;
    for (int i = -1; i < numIterations; i++)
    {
        if (i == 0) 
        {
            CUDA_SAFE_CALL(cudaThreadSynchronize());
        }

        // Each thread will generate N_PER_RNG randoms
        RandomGPU<<<GRIDSZ, BLOCKSZ>>>(d_Rand,
                                       N_PER_RNG);
    }

#if 0
    float *h_Rand  = (float *)malloc(RAND_N * sizeof(float));

    CUDA_SAFE_CALL(cudaMemcpy(h_Rand, d_Rand, RAND_N * sizeof(float), cudaMemcpyDeviceToHost));

    for (unsigned int i = 0; i < RAND_N; i++)
    {
        printf("%+2.8f ", h_Rand[i]);
        if (i % 10 == 0)
        {
            printf("\n");
        }
    }
    printf("\n");

    free(h_Rand);
#endif
}

void initIFS(void)
{
    // Initialize variations. First decide how many.
    numFunctions = 3 + randInt(1, 8);

    // Init the coefficients of the variations
    for (int i = 0; i < numFunctions; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            // Coefficients lie in [-1, 1]
            coeffMatrix[i][j] = randFloat(-1, 1);
        }

        probability[i] = 0.1 + randFloat(0, 1);

        functionType[i] = weightedRandom(variationWeights,
                                         numFunctions,
                                         randFloat(0, 1));

        variationColors[i][RED] = 0;
        variationColors[i][GREEN] = 0;
        variationColors[i][BLUE] = 0;
    }

    CUDA_SAFE_CALL(cudaMalloc((void **)&d_Probability,
                               numFunctions * sizeof(float)));
    CUDA_SAFE_CALL(cudaMalloc((void **)&d_Coefficients,
                               numFunctions * 6 * sizeof(float)));

    CUDA_SAFE_CALL(cudaMemcpy(d_Probability, probability,
                               numFunctions * sizeof(float),
                               cudaMemcpyHostToDevice));
    CUDA_SAFE_CALL(cudaMemcpy(d_Coefficients, coeffMatrix,
                               numFunctions * 6 * sizeof(float),
                               cudaMemcpyHostToDevice));
}

int randInt(int min, int max)
{
    // overflows int without typecast to RAND_MAX
    float fraction = (float)rand() / ((float)RAND_MAX + 1);
    return (int)(fraction * (max - min) + min);
}

float
randFloat(float min, float max)
{
    // overflows int without typecast to RAND_MAX
    float fraction = (float)rand() / ((float)RAND_MAX + 1);
    return (fraction * (max - min) + min);
}

// selects a random index inside "weights", and probability is based on the
// weights themselves. "count" is both the number of weights and the index
// range. randomNum is a random number in the [0, 1] interval.
__host__ __device__ int 
weightedRandom(float *weights, int count, float randomNum)
{
    float cumulativeWeight[MAX_FUNCTIONS];

    cumulativeWeight[0] = 0;

    for (int i = 1; i <= count ; i++)
    {
        cumulativeWeight[i] += cumulativeWeight[i-1] + weights[i-1];
    }

    float weightedRandomNum = randomNum * cumulativeWeight[count];

    int index;
    for (index = 0; cumulativeWeight[index] < weightedRandomNum; ++index);

    return index-1;
}

// Update image buffer with processed coordinates
// As input this takes an array of coordinates that have been processed by the
// IFS. We need to look at each coordinate, and then update the
// colour/brightness etc of the corresponding coordinate in the image

